#include <cv_bridge/cv_bridge.h>
#include <rmw/qos_profiles.h>
#include <tf2/LinearMath/Matrix3x3.h>
#include <tf2/convert.h>

#include <ament_index_cpp/get_package_share_directory.hpp>
#include <image_transport/image_transport.hpp>
#include <opencv2/core.hpp>
#include <opencv2/imgproc.hpp>
#include <rclcpp/duration.hpp>
#include <rclcpp/qos.hpp>
#include <tf2_geometry_msgs/tf2_geometry_msgs.hpp>

// STD
#include <algorithm>
#include <map>
#include <memory>
#include <string>
#include <vector>

#include "armor_detector/armor.hpp"
#include "armor_detector/detector_node.hpp"

namespace rm_auto_aim
{
  ArmorDetectorNode::ArmorDetectorNode(const rclcpp::NodeOptions & options)
  : Node("armor_detector", options)
  {
    RCLCPP_INFO(this->get_logger(), "Starting DetectorNode!");

    // Detector(初始化检测器)
    detector_ = initDetector();

    // Armors Publisher(发布检测到的装甲板信息)
    armors_pub_ = this->create_publisher<auto_aim_interfaces::msg::Armors>(
      "/detector/armors", rclcpp::SensorDataQoS());

    // Visualization Marker Publisher(初始化可视化标记)
    // See http://wiki.ros.org/rviz/DisplayTypes/Marker
    armor_marker_.ns = "armors";  // 命名空间来区分不同的标记组
    armor_marker_.action = visualization_msgs::msg::Marker::ADD;  // 添加或更新标记
    armor_marker_.type = visualization_msgs::msg::Marker::CUBE;   // 表示立方体
    armor_marker_.scale.x = 0.05; // 立方体在x轴方向的尺寸
    armor_marker_.scale.z = 0.125;  // 立方体在z轴方向的尺寸
    armor_marker_.color.a = 1.0;  // 透明度(1.0为完全不透明)
    armor_marker_.color.g = 0.5;  // 绿色分量
    armor_marker_.color.b = 1.0;  // 蓝色分量
    armor_marker_.lifetime = rclcpp::Duration::from_seconds(0.1); // 标记的存在时间是0.1秒

    text_marker_.ns = "classification";
    text_marker_.action = visualization_msgs::msg::Marker::ADD;
    text_marker_.type = visualization_msgs::msg::Marker::TEXT_VIEW_FACING;  // 表示始终面向摄像头的文本
    text_marker_.scale.z = 0.1; // 文本的高度
    text_marker_.color.a = 1.0;  
    text_marker_.color.r = 1.0; // 红色分量
    text_marker_.color.g = 1.0;
    text_marker_.color.b = 1.0;
    text_marker_.lifetime = rclcpp::Duration::from_seconds(0.1);

    marker_pub_ =
      this->create_publisher<visualization_msgs::msg::MarkerArray>("/detector/marker", 10);

    // Debug Publishers(调试模式)
    debug_ = this->declare_parameter("debug", true); // 默认值是false,改成了true
    if (debug_) 
    {
      createDebugPublishers();
    }

    // Task subscriber(任务订阅)
    is_aim_task_ = true;
    task_sub_ = this->create_subscription<std_msgs::msg::String>(
      "/task_mode", 10, std::bind(&ArmorDetectorNode::taskCallback, this, std::placeholders::_1));

    // Debug param change moniter(调试参数监听)
    debug_param_sub_ = std::make_shared<rclcpp::ParameterEventHandler>(this);
    debug_cb_handle_ =
      debug_param_sub_->add_parameter_callback("debug", [this](const rclcpp::Parameter & p) {
        debug_ = p.as_bool();
        debug_ ? createDebugPublishers() : destroyDebugPublishers();
      });

    // 相机信息订阅
    cam_info_sub_ = this->create_subscription<sensor_msgs::msg::CameraInfo>(
      "/camera_info", rclcpp::SensorDataQoS(),
      [this](sensor_msgs::msg::CameraInfo::ConstSharedPtr camera_info) {
        cam_center_ = cv::Point2f(camera_info->k[2], camera_info->k[5]);
        cam_info_ = std::make_shared<sensor_msgs::msg::CameraInfo>(*camera_info);
        pnp_solver_ = std::make_unique<PnPSolver>(camera_info->k, camera_info->d);
        cam_info_sub_.reset();
      });

    // 图像订阅
    img_sub_ = this->create_subscription<sensor_msgs::msg::Image>(
      "/image_raw", rclcpp::SensorDataQoS(),
      std::bind(&ArmorDetectorNode::imageCallback, this, std::placeholders::_1));
  }

  void ArmorDetectorNode::taskCallback(const std_msgs::msg::String::SharedPtr task_msg)
  {
    std::string task_mode = task_msg->data;
    if (task_mode == "aim") 
    {
      is_aim_task_ = true;
    } 
    else 
    {
      is_aim_task_ = false;
    }
  }

  void ArmorDetectorNode::imageCallback(const sensor_msgs::msg::Image::ConstSharedPtr img_msg)
  {
    auto armors = detectArmors(img_msg);

    if (pnp_solver_ != nullptr && is_aim_task_) 
    {
      // RCLCPP_INFO(get_logger(), "OK!!!");
      // 初始化消息和标记
      armors_msg_.header = armor_marker_.header = text_marker_.header = img_msg->header;
      armors_msg_.armors.clear();
      marker_array_.markers.clear();
      armor_marker_.id = 0;
      text_marker_.id = 0;

      // 遍历检测到的装甲板
      auto_aim_interfaces::msg::Armor armor_msg;
      for (const auto & armor : armors) 
      {
        cv::Mat rvec, tvec;
        bool success = pnp_solver_->solvePnP(armor, rvec, tvec);  // 使用PnP求解器计算3D位置和姿态(rvec: 旋转向量, tvec: 平移向量)
        if (success)  // 填充装甲板消息
        {
          // Fill basic info
          armor_msg.type = ARMOR_TYPE_STR[static_cast<int>(armor.type)];
          armor_msg.number = armor.number;

          // Fill pose
          armor_msg.pose.position.x = tvec.at<double>(0);
          armor_msg.pose.position.y = tvec.at<double>(1);
          armor_msg.pose.position.z = tvec.at<double>(2);
          // rvec to 3x3 rotation matrix
          cv::Mat rotation_matrix;
          cv::Rodrigues(rvec, rotation_matrix);
          // rotation matrix to quaternion(旋转向量转换为四元数)
          tf2::Matrix3x3 tf2_rotation_matrix(
            rotation_matrix.at<double>(0, 0), rotation_matrix.at<double>(0, 1),
            rotation_matrix.at<double>(0, 2), rotation_matrix.at<double>(1, 0),
            rotation_matrix.at<double>(1, 1), rotation_matrix.at<double>(1, 2),
            rotation_matrix.at<double>(2, 0), rotation_matrix.at<double>(2, 1),
            rotation_matrix.at<double>(2, 2));
          tf2::Quaternion tf2_q;
          tf2_rotation_matrix.getRotation(tf2_q);
          armor_msg.pose.orientation = tf2::toMsg(tf2_q);

          // Fill the distance to image center(计算装甲板到图像中心的距离)
          armor_msg.distance_to_image_center = pnp_solver_->calculateDistanceToCenter(armor.center);

          // Fill keypoints(填充关键点)
          armor_msg.kpts.clear();
          for (const auto & pt :
              {armor.left_light.top, armor.left_light.bottom, armor.right_light.bottom,
                armor.right_light.top}) 
          {
            geometry_msgs::msg::Point point;
            point.x = pt.x;
            point.y = pt.y;
            armor_msg.kpts.emplace_back(point);
          }

          // Fill the markers(填充可视化标记)
          armor_marker_.id++;
          armor_marker_.scale.y = armor.type == ArmorType::SMALL ? 0.135 : 0.23;
          armor_marker_.pose = armor_msg.pose;
          text_marker_.id++;
          text_marker_.pose.position = armor_msg.pose.position;
          text_marker_.pose.position.y -= 0.1;
          text_marker_.text = armor.classfication_result;
          // 将装甲板消息和标记添加到消息列表和标记数组中
          armors_msg_.armors.emplace_back(armor_msg);
          marker_array_.markers.emplace_back(armor_marker_);
          marker_array_.markers.emplace_back(text_marker_);
        } 
        else 
        {
          RCLCPP_WARN(this->get_logger(), "PnP failed!");
        }
      }

      // Publishing detected armors
      armors_pub_->publish(armors_msg_);

      // Publishing marker
      publishMarkers();
    }
  }

  std::unique_ptr<Detector> ArmorDetectorNode::initDetector()
  {
    rcl_interfaces::msg::ParameterDescriptor param_desc;
    param_desc.integer_range.resize(1);
    param_desc.integer_range[0].step = 1;
    param_desc.integer_range[0].from_value = 0;
    param_desc.integer_range[0].to_value = 255;
    int binary_thres = declare_parameter("binary_thres", 70, param_desc);  // 二值化阈值(0~255) 160

    // 检测颜色
    param_desc.description = "0-RED, 1-BLUE";
    param_desc.integer_range[0].from_value = 0;
    param_desc.integer_range[0].to_value = 1;
    auto detect_color = declare_parameter("detect_color", BLUE, param_desc);

    // 灯条参数
    Detector::LightParams l_params = {
      .min_ratio = declare_parameter("light.min_ratio", 0.1), // 最小长宽比
      .max_ratio = declare_parameter("light.max_ratio", 0.4), // 最大长宽比
      .max_angle = declare_parameter("light.max_angle", 35.0), // 最大倾斜角
      .min_fill_ratio = declare_parameter("light.min_fill_ratio", 0.8), // 最小填充率
    };

    // 装甲板参数
    Detector::ArmorParams a_params = {
      .min_light_ratio = declare_parameter("armor.min_light_ratio", 0.7),                     // 装甲板中灯条的最小比例
      .min_small_center_distance = declare_parameter("armor.min_small_center_distance", 0.8), // 小装甲板灯条中心的最小距离
      .max_small_center_distance = declare_parameter("armor.max_small_center_distance", 3.2), // 小装甲板灯条中心的最大距离
      .min_large_center_distance = declare_parameter("armor.min_large_center_distance", 3.2), // 大装甲板灯条中心的最小距离
      .max_large_center_distance = declare_parameter("armor.max_large_center_distance", 5.5), // 大装甲板灯条中心的最大距离
      .max_angle = declare_parameter("armor.max_angle", 35.0)                                 // 装甲板的最大倾斜角度
      };

    // 加载参数初始化
    auto detector = std::make_unique<Detector>(binary_thres, detect_color, l_params, a_params);

    // Init classifier(初始化数字分类器)
    auto pkg_path = ament_index_cpp::get_package_share_directory("armor_detector"); // 获取路径
    auto model_path = pkg_path + "/model/mlp.onnx"; // 数字分类器的模型路径
    auto label_path = pkg_path + "/model/label.txt";  // 数字分类器的标签路径
    double threshold = this->declare_parameter("classifier_threshold", 0.7);  // 分类器的置信度阈值
    std::vector<std::string> ignore_classes =
      this->declare_parameter("ignore_classes", std::vector<std::string>{"negative"});  // 忽略的分类标签
    detector->classifier =
      std::make_unique<NumberClassifier>(model_path, label_path, threshold, ignore_classes);

    return detector;
  }

  std::vector<Armor> ArmorDetectorNode::detectArmors(
    const sensor_msgs::msg::Image::ConstSharedPtr & img_msg)
  {
    // Convert ROS img to cv::Mat (将ros图像转换成opencv可以处理的图像)
    auto img = cv_bridge::toCvShare(img_msg, "rgb8")->image;

    // Update params(更新参数)
    detector_->binary_thres = get_parameter("binary_thres").as_int();
    detector_->detect_color = get_parameter("detect_color").as_int();
    detector_->classifier->threshold = get_parameter("classifier_threshold").as_double();

    auto armors = detector_->detect(img);
    // RCLCPP_INFO(get_logger(), "OK");

    // 计算检测延迟
    auto final_time = this->now();  // 当前时间
    auto latency = (final_time - img_msg->header.stamp).seconds() * 1000; // 从图像时间戳到当前时间的差值
    RCLCPP_DEBUG_STREAM(this->get_logger(), "Latency: " << latency << "ms");  // 记录延迟日志

    // Publish debug info(发布调试信息)
    if (debug_) 
    {
      binary_img_pub_.publish(
        cv_bridge::CvImage(img_msg->header, "mono8", detector_->binary_img).toImageMsg());  // 发布二值化图像

      // Sort lights and armors data by x coordinate  // 按x坐标对灯条和装甲板数据进行排序
      std::sort(
        detector_->debug_lights.data.begin(), detector_->debug_lights.data.end(),
        [](const auto & l1, const auto & l2) { return l1.center_x < l2.center_x; });
      std::sort(
        detector_->debug_armors.data.begin(), detector_->debug_armors.data.end(),
        [](const auto & a1, const auto & a2) { return a1.center_x < a2.center_x; });

      // 发布灯条和装甲板数据
      lights_data_pub_->publish(detector_->debug_lights);
      armors_data_pub_->publish(detector_->debug_armors);

      if (!armors.empty()) 
      {
        auto all_num_img = detector_->getAllNumbersImage();
        number_img_pub_.publish(
          *cv_bridge::CvImage(img_msg->header, "mono8", all_num_img).toImageMsg()); // 发布数字图像
      }

      // 在图像上绘制检测结果
      detector_->drawResults(img);
      // Draw camera center(绘制相机中心点)
      cv::circle(img, cam_center_, 5, cv::Scalar(255, 0, 0), 2);
      // Draw latency(在图像上显示检测延迟)
      std::stringstream latency_ss;
      latency_ss << "Latency: " << std::fixed << std::setprecision(2) << latency << "ms";
      auto latency_s = latency_ss.str();
      cv::putText(
        img, latency_s, cv::Point(10, 30), cv::FONT_HERSHEY_SIMPLEX, 1.0, cv::Scalar(0, 255, 0), 2);
      result_img_pub_.publish(cv_bridge::CvImage(img_msg->header, "rgb8", img).toImageMsg()); // 发布结果图像
    }

    return armors;
  }

  void ArmorDetectorNode::createDebugPublishers()
  {
    lights_data_pub_ =
      this->create_publisher<auto_aim_interfaces::msg::DebugLights>("/detector/debug_lights", 10);
    armors_data_pub_ =
      this->create_publisher<auto_aim_interfaces::msg::DebugArmors>("/detector/debug_armors", 10);

    binary_img_pub_ = image_transport::create_publisher(this, "/detector/binary_img");
    number_img_pub_ = image_transport::create_publisher(this, "/detector/number_img");
    result_img_pub_ = image_transport::create_publisher(this, "/detector/result_img");
  }

  void ArmorDetectorNode::destroyDebugPublishers()
  {
    lights_data_pub_.reset();
    armors_data_pub_.reset();

    binary_img_pub_.shutdown();
    number_img_pub_.shutdown();
    result_img_pub_.shutdown();
  }

  void ArmorDetectorNode::publishMarkers()
  {
    using Marker = visualization_msgs::msg::Marker;
    armor_marker_.action = armors_msg_.armors.empty() ? Marker::DELETE : Marker::ADD;
    marker_array_.markers.emplace_back(armor_marker_);
    marker_pub_->publish(marker_array_);
  }

}  // namespace rm_auto_aim

#include "rclcpp_components/register_node_macro.hpp"

// Register the component with class_loader.
// This acts as a sort of entry point, allowing the component to be discoverable when its library
// is being loaded into a running process.
RCLCPP_COMPONENTS_REGISTER_NODE(rm_auto_aim::ArmorDetectorNode)
